# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: chensili, leon

import datetime
import os

from absl import app, flags
import glob

from coin.base.datetime_util import (
    convert_string_to_dates,
    iterate_date,
)
from coin.strategy.mm.subscription import SubscriptionRequest
from coin.exchange.kr_rest.product.product_impl import get_exchange_product_list
from coin.support.feed_tool.feed_stats.app.feed_stats_runner import EXCHANGE_LIST
from coin.strategy.mm.feed import get_flow_book_builder_cls

FLAGS = flags.FLAGS

feed_archive_colo_configs = [
    {
        'colo': 'feed-01.ap-northeast-2.aws',
        'exchanges': [
            'Okex',
            'Binance',
            'Bitfinex',
            'Bitmex',
            'Bithumb',
            'Upbit',
            'Korbit',
            'Coinone',
            'Gopax'
        ],
        'coin_group': 'all'
    },
    {
        'colo': 'feed-04.ap-northeast-1.aws',
        'exchanges': [
            'Kraken',
            'Bitstamp',
            'Bitflyer',
            'Quoinex',
            'Hitbtc',
            'Bequant',
            'Gdax',
            'Deribit',
            'Bybit'
        ],
        'coin_group': 'all'
    },
]

machine_to_exchanges = {config['colo']: config['exchanges'] for config in feed_archive_colo_configs}


def _check_feed_exist(machine, date, root_dir=None):
  date_str = date.strftime('%Y%m%d')
  root_dir = root_dir or '/remote/iosg/coin/data/flow'

  search_topics = set()
  for market_type, exchange, api_version in EXCHANGE_LIST:
    if exchange not in machine_to_exchanges[machine]:
      continue
    products = get_exchange_product_list(market_type, exchange, api_version, date)
    sub_req = SubscriptionRequest(market_type, exchange, api_version)
    flow_book_builder_cls = get_flow_book_builder_cls(sub_req)
    for product in products:
      for worker_id in ['1', '2']:
        topics = flow_book_builder_cls.get_topic_string(product, worker_id)
        for topic in topics:
          if topic in search_topics:
            continue
          else:
            search_topics.add(topic)

          search_str = os.path.join(root_dir, machine, date_str, topic) + '*'
          files = glob.glob(search_str)
          if not files:
            print("Topic %s not downloaded." % topic)


def check_download_status(root_dir, start_date: datetime.date, end_date: datetime.date):
  machines = machine_to_exchanges.keys()
  for machine in machines:
    print('\nMachine: %s' % machine)
    machine_dir = os.path.join(root_dir, machine)
    if not os.path.isdir(machine_dir):
      print('Dir not exist: %s' % machine_dir)
      continue
    for date in iterate_date(start_date, end_date):
      date_str = date.strftime('%Y%m%d')
      print('***** Trading_date: %s' % date_str)
      date_dir = os.path.join(machine_dir, date_str)
      if not os.path.isdir(date_dir):
        print('Dir not exist: %s' % date_dir)
        continue

      _check_feed_exist(machine, date, root_dir)


def main(_):
  start_date = convert_string_to_dates(FLAGS.start_date)[0]
  end_date = convert_string_to_dates(FLAGS.end_date)[0]
  check_download_status(FLAGS.root_dir, start_date, end_date)


if __name__ == '__main__':
  flags.DEFINE_string('root_dir', '/remote/iosg/coin/data/flow', '')

  flags.DEFINE_string('start_date', None, 'yyyymmdd')

  flags.DEFINE_string('end_date', None, 'yyyymmdd. Feed at end_date is not checked.')

  app.run(main)
